import time

import numpy as np
import pandas as pd
from pymoo.core.problem import ElementwiseProblem
from pymoo.algorithms.moo.nsga2 import NSGA2
from pymoo.operators.crossover.sbx import SBX
from pymoo.operators.mutation.pm import PM
from pymoo.operators.sampling.rnd import FloatRandomSampling
from pymoo.termination import get_termination
from pymoo.optimize import minimize
import matplotlib.pyplot as plt

# ====================== 基础参数设置 ======================
T = 144  # 3天×24小时×2时段/小时
dt = 0.5 * 3600  # 时段长度（秒）

def storage_sg_func(v_array):
    """
    向量化计算存储函数值（支持ndarray输入）

    参数:
    v_array (ndarray): 待计算的容积数组（单位：万m³）

    返回:
    ndarray: 计算结果数组（超出区间的值为np.nan）
    """
    # 转换为浮点数组避免整数运算问题
    v = v_array.astype(np.float64)

    # 定义区间条件
    cond1 = np.logical_and(v > -np.inf, v <= 2.58)  # 区间 [0, 2.58]
    cond2 = np.logical_and(v > 2.58, v <= np.inf)  # 区间 (2.58, 311.02]

    # 计算各区间的多项式值
    poly1 = -0.4068 * (v ** 3) + 0.3595 * (v ** 2) + 0.5759 * v + 4.5012
    poly2 = 2e-07 * (v ** 3) - 0.0001 * (v ** 2) + 0.0472 * v + 5.1436

    # 根据条件选择结果（超出区间返回np.nan）
    result = np.select([cond1, cond2], [poly1, poly2], default=np.nan)

    return result

def storage_sg_inv_func(h_array):
    """
    向量化计算水位-库容反函数值（支持ndarray输入）

    参数:
    h_array (ndarray): 待计算的水位数组（单位：m）

    返回:
    ndarray: 计算结果数组（超出区间的值为np.nan）
    """
    # 转换为浮点数组避免整数运算问题
    h = h_array.astype(np.float64)

    # 定义区间条件
    cond1 = np.logical_and(h >= 4.5, h <= 5.11)  # 区间 [4.5, 5.11]
    cond2 = np.logical_and(h > 5.11, h <= 13)  # 区间 (5.11, 13]

    # 计算各区间的多项式值
    poly1 = 4.0878 * (h ** 3) - 57.49 * (h ** 2) + 270.92 * h - 427.45
    poly2 = -0.3646 * (h ** 3) + 12.367 * (h ** 2) - 89.848 * h + 190.18

    # 根据条件选择结果（超出区间返回np.nan）
    result = np.select([cond1, cond2], [poly1, poly2], default=np.nan)

    return result
# 水库参数（保持不变）
reservoir_params = {
    'A': {
        'storage_func': lambda V: -0.0000004 * V ** 2 + 0.0052 * V + 14.783,  # 库容-水位函数 V(万m³)->Z(m)
        'storage_inv_func': lambda H: 13.767 * H ** 2 - 362.59 * H + 2868.2,  # 水位-库容反函数 Z(m)->V(万m³)
        'Q_min': 0,  # 最小出库流量(m³/s)
        'Q_max': 270,  # 最大出库流量(m³/s)
        'H_initial': 26.6,
        'V_initial': 3348.4  # 初始库容(万m³) 深圳水库
    },
    'B': {
        'storage_func': 0,
        'storage_inv_func': 0,
        'Q_min': 0,
        'Q_max': 300,
        'H_initial': 11.73,
        'V_initial': 240.61  # 笋岗滞洪区
    }
}
gates_params = {
    'gate_SCHK': {
        'H_initial': 2.5,
        'Q_to_H_func': lambda Q: 0.1568 * (Q ** 0.5494)  # 幂函数：H = a*Q^b，其中a=0.1568，b=0.5494 断面流量-水位关系
    },
    'gate_WJD': {
        'H_initial': 2.5,
        'Q_to_H_func': lambda Q: 0.1361 * (Q ** 0.5685)  #  断面流量-水位关系
    },
    'gate_LHZ': {
        'H_initial': 2.5,
        'Q_to_H_func': lambda Q: 0.0907 * (Q ** 0.6065)  # 幂函数：H = a*Q^b，其中a=0.1568，b=0.5494 断面流量-水位关系
    },
    'gate_BJHK': {
        'H_initial': 2.5,
        'Q_to_H_func': lambda Q: 0.3943 * (Q ** 0.3917)
    },
    'gate_LDC': {
        'H_initial': 2.5,
        'Q_to_H_func': lambda Q: 0.0675 * (Q ** 0.5805)
    },
    'gate_SBMT': {
        'H_initial': 2.5,
        'Q_to_H_func': lambda Q: 0.0505 * (Q ** 0.6086)
    },
    'gate_SZHK': {
        'H_initial': 2.5,
        'Q_to_H_func': lambda Q: 0.0011 * Q + 0.7779
    }

}
# 水闸流达时间（分钟）→ 转换为时段数（半小时为一时段）
travel_time = {
    'A到三岔河口': 6,  # A水库到三岔河口
    '三岔河口到文锦渡': 8.1,  # 三岔河口到文锦渡
    '文锦渡到罗湖站': 25,  # 文锦渡到罗湖站
    '罗湖站到布吉河口': 13,  # 罗湖站到布吉河口
    'B到布吉河口': 12,  # B水库到布吉河口
    '布吉河口、罗湖站到鹿丹村': 38,  # 布吉河口到鹿丹村
    '鹿丹村到上步码头': 10,  # 鹿丹村到上步码头
    '上步码头到深圳河口': 6  # 上步码头到深圳河口
}
# 转换为时段数（向上取整）
travel_time_period = {k: int(np.ceil(v / 30)) for k, v in travel_time.items()}

def load_inflow_data(excel_path):
    try:
        # 从Excel读取数据，假设Excel文件包含'A_inflow'和'B_inflow'两列
        df = pd.read_excel(excel_path)

        # 检查数据行数是否为时段数
        if len(df) != T:
            raise ValueError(f"数据行数应为{T}行，当前读取到{len(df)}行")

        # 检查是否包含所需列
        required_columns = ['Date Time','笋岗滞洪区入流', '深圳水库入流']
        if not set(required_columns).issubset(df.columns):
            missing = [col for col in required_columns if col not in df.columns]
            raise KeyError(f"Excel文件缺少必要列: {missing}")

        return df[required_columns]

    except FileNotFoundError:
        print(f"错误：未找到文件 {excel_path}")
        return None
    except Exception as e:
        print(f"读取数据时发生错误: {str(e)}")
        return None


# 200年一遇洪水过程 3天
excel_file = "inflow_data_200a.xlsx"
inflow_data = load_inflow_data(excel_file)
A_lateral_inflow = np.zeros(T)  # 深圳水库旁侧入流（来水）
B_lateral_inflow = np.zeros(T)  # 笋岗滞洪区旁侧入流（来水）
if inflow_data is not None:
    print("成功读取入库流量数据：")
    print(inflow_data.head())  # 打印前5行验证

def compute_flow(upstream_flow, delay, total_time):
    """
    通用流量计算函数
    :param upstream_flow: 上游流量数组
    :param delay: 延迟时段
    :param total_time: 总时间步长T
    :return: 当前节点流量数组
    """
    return np.array([
        upstream_flow[t - delay] if t >= delay else 0
        for t in range(total_time)
    ])

class MyProblem(ElementwiseProblem):

    def __init__(self, **kwargs):
        super().__init__(n_var=2 * T,
                         n_obj=2,
                         n_ieq_constr=6 * T,
                         xl=np.array([reservoir_params['A']['Q_min']] * T + [reservoir_params['B']['Q_min']] * T),
                         # 出库下限
                         xu=np.array([reservoir_params['A']['Q_max']] * T + [reservoir_params['B']['Q_max']] * T),
                         **kwargs)  # 在问题定义时可以设置并行计算的参数

        self.flow_processes = []  # 保存水库、断面的流量过程

    def _evaluate(self, x, out, *args, **kwargs):
        f = []  # 目标函数值（洪峰流量）
        g = []  # 约束条件
        # 分解决策变量
        Q_A = x[:T]  # 深圳水库A各时段出库流量(m³/s)
        Q_B = x[T:]  # 笋岗滞洪区B各时段出库流量(m³/s)
        # ---------------------- 水库调蓄计算 ----------------------
        # 深圳水库A水量平衡
        V_A = np.zeros(T + 1)  # 库容序列 万m³
        H_A = np.zeros(T + 1)  # 深圳水库水位 m
        V_A[0] = reservoir_params['A']['V_initial']  # 深圳水库初始流量
        H_A[0] = reservoir_params['A']['H_initial']  # 深圳水库初始水位

        # 笋岗滞洪区B水量平衡
        V_B = np.zeros(T + 1)  # 笋岗滞洪区库容序列（万m³）
        H_B = np.zeros(T + 1)  # 笋岗滞洪区水位 m
        V_B[0] = reservoir_params['B']['V_initial']  # 笋岗滞洪区初始水位
        H_B[0] = reservoir_params['B']['H_initial']  # 笋岗滞洪区初始水位
        for t in range(T):
            inflow = inflow_data['深圳水库入流'].iloc[t]  # 深圳水库入库流量(m³/s)
            inflow1 = inflow_data['笋岗滞洪区入流'].iloc[t]  # 笋岗滞洪区入库流量(m³/s)
            V_A[t + 1] = V_A[t] + (inflow - Q_A[t]) * dt / 10000  # 转换为万m³
            H_A[t + 1] = reservoir_params['A']['storage_func'](V_A[t + 1])

            # if V_B[t] < 300:
            #     Q_B[t] = 0
            #     V_B[t + 1] = V_B[t] + (inflow1 - Q_B[t]) * dt / 10000
            #     H_B[t + 1] = storage_sg_func(V_B[t + 1])
            # if inflow1 > Q_B[t] and inflow1 >300 :
            #     Q_B[t] = inflow1
            V_B[t + 1] = V_B[t] + (inflow1 - Q_B[t]) * dt / 10000
            H_B[t + 1] = storage_sg_func(V_B[t + 1])

        # ---------------------- 下游水闸流量计算 ----------------------
        # 1. 计算布罗湖站流量（A和B的汇流）
        # A到罗湖站的延迟时段路径：A→三岔河口→文锦渡→罗湖站
        delay_dict = {
            'A到罗湖站': (travel_time_period['A到三岔河口'] +
                          travel_time_period['三岔河口到文锦渡'] +
                          travel_time_period['文锦渡到罗湖站']),
            'B到布吉河口': travel_time_period['B到布吉河口'],
            '布吉罗湖到鹿丹村': travel_time_period['布吉河口、罗湖站到鹿丹村'],
            '鹿丹村到上步码头': travel_time_period['鹿丹村到上步码头'],
            '上步到深圳河口': travel_time_period['上步码头到深圳河口']
        }

        # 1. 计算罗湖站流量（深圳水库→罗湖）
        Q_LHZ = compute_flow(Q_A, delay_dict['A到罗湖站'], T)
        # 罗湖水位
        H_LHZ = np.zeros(T)
        for t in range(T):
            if t == 0:
                H_LHZ[t] = gates_params['gate_LHZ']['H_initial']
            elif Q_LHZ[t] == 0:
                H_LHZ[t] = H_LHZ[t - 1]
            else:
                H_LHZ[t] = gates_params['gate_LHZ']['Q_to_H_func'](Q_LHZ[t])
        # 2.笋岗滞洪区到布吉河口的流量
        Q_BJHK = compute_flow(Q_B, delay_dict['B到布吉河口'], T)
        # TODO:布吉河水位
        H_BJHK = np.zeros(T)
        for t in range(T):
            if t == 0:
                H_BJHK[t] = gates_params['gate_BJHK']['H_initial']
            elif Q_BJHK[t] == 0:
                H_BJHK[t] = H_BJHK[t - 1]
            else:
                H_BJHK[t] = gates_params['gate_BJHK']['Q_to_H_func'](Q_BJHK[t])
        # 3.计算鹿丹村流量（布吉河口+ B→布吉河口 共同延迟到鹿丹村）
        Q_LDC = compute_flow(Q_LHZ + Q_BJHK, delay_dict['布吉罗湖到鹿丹村'], T)
        # TODO:鹿丹村水位
        # 4.计算上步码头流量（鹿丹村→上步码头）
        Q_SBMT = compute_flow(Q_LDC, delay_dict['鹿丹村到上步码头'], T)
        # TODO:上步码头水位
        # 5. 计算深圳河口流量（上步码头→深圳河口）
        Q_SZHK = compute_flow(Q_SBMT, delay_dict['上步到深圳河口'], T)
        # TODO:深圳河口水位

        # ---------------------- 目标函数：下游最大洪峰流量 ----------------------
        # 各水闸的洪峰流量（过程最大值）

        peak_LHZ = np.max(Q_LHZ)
        peak_BJHK = np.max(Q_BJHK)
        peak_LDC = np.max(Q_LDC)
        peak_SBMT = np.max(Q_SBMT)
        peak_SZHK = np.max(Q_SZHK)

        # 下游整体最大洪峰（取所有水闸的最大值）
        total_peak = np.max([peak_LHZ, peak_BJHK, peak_LDC, peak_SBMT, peak_SZHK])

        # total_v = 0
        # for t in range(T):
        #     total_v += (V_B[t+1] - V_B[t]) ** 2

        # 水位波动最小
        total_h = 0
        for t in range(T):
            total_h += (0
                        # (H_LHZ[t] - gates_params['gate_LHZ']['H_initial']) ** 2
                        + (H_BJHK[t] - gates_params['gate_BJHK']['H_initial']) ** 2
                        )

        total_v = 0
        for t in range(T):
            total_v += (
                    (V_A[t + 1] - reservoir_params['A']['V_initial']) ** 2
                    + (V_B[t + 1] - reservoir_params['B']['V_initial']) ** 2
            )


        # ---------------------- 约束条件 ----------------------
        # 1. 深圳水库时段末库容非负（V(t+1) ≥ 0）
        constr_vA = -V_A[1:]  # -V ≤ 0 → V ≥ 0

        # 2. 笋岗滞洪区时段末库容非负
        constr_vB = -V_B[1:]
        constr_vB1 = 311 - V_B[1:]
        # 3. 深圳水库水位不超设计上限
        Z_A = reservoir_params['A']['storage_func'](V_A[1:])  # 时段初水位
        constr_zA = Z_A - 28.83  # Z - Z_max ≤ 0

        constr_zA1 = 24.3 - Z_A
        # 4. 笋岗滞洪区水位不超设计上限
        constr_zB = H_B[1:] - 13
        # 5. 笋岗滞洪区水位不超设计下限
        constr_zB1 = 6 - H_B[1:]

        # 6. 罗湖站断面水位不超设计上限
        constr_h_LHZ = H_LHZ - 5.816  # 上限 5.816
        constr_h_LHZ1 = 0 - H_LHZ  # 下限0

        constr_h_BJHK = H_BJHK - 5.15  # 上限 5.15
        constr_h_BJHK1 = 0 - H_BJHK  # 下限0
        # 合并所有约束（g ≤ 0 表示可行）
        g.append(np.concatenate([constr_vA,
                                 constr_vB,
                                 # constr_vB1,
                                 constr_zA,
                                 constr_zA1,
                                 # constr_zB,
                                 # constr_zB1,
                                 constr_h_LHZ,
                                 constr_h_LHZ1,
                                 # constr_h_BJHK,
                                 # constr_h_BJHK1
                                 ])
                 )
        # 保存当前解的流量过程
        self.flow_processes.append({
            '深圳水库出库流量': Q_A,
            '笋岗滞洪区出库流量': Q_B,
            '深圳水库水位': H_A[1:],
            '深圳水库库容': V_A[1:],
            '笋岗滞洪区水位': H_B[1:],
            '笋岗滞洪区库容': V_B[1:],
            '罗湖站流量': Q_LHZ,
            '罗湖站水位': H_LHZ,
            '布吉河口流量': Q_BJHK,
            '布吉河口水位': H_BJHK,
        })
        out["F"] = np.column_stack([total_v, total_peak])
        out["G"] = g
def optimization_callback(algorithm):
    """回调函数：输出每代优化进度"""
    gen = algorithm.n_gen
    current_min_peak = np.min(algorithm.pop.get("F"))
    print(f"第 {gen:3d} 代 | 当前水位违反值: {current_min_peak:.2f} m")

if __name__ == "__main__":
    start_time = time.perf_counter()
problem = MyProblem()

algorithm = NSGA2(
    pop_size=41,  # 种群数量
    n_offsprings=10,  # 每一代中从父代个体中产生的新个体的数量
    sampling=FloatRandomSampling(),  # 用于生成初始种群的采样算子。
    crossover=SBX(prob=0.9, eta=15),  # 交叉算法，模拟二进制
    # 交叉概率，分布指数，变异指数，上下限
    mutation=PM(eta=20),  # 变异算子，默认为多项式变异
    eliminate_duplicates=True,  # 剔除重复项
)

termination = get_termination("n_gen", 500)

res = minimize(problem,
               algorithm,
               termination,
               seed=1,  # seed 相同就能保证在任何地方执行的结果相同
               save_history=True,# 保存历史
               callback=optimization_callback,
               verbose=False)  # 输出优化过程

X = res.X  # 结果中的自变量
F = res.F  # 结果中的目标量
print(f"最终种群目标值: {res.F}")
xl, xu = problem.bounds()

# 输出结果
# print(f"\n优化完成,最终最小水位违反值: {res.F[0]:.2f} m")

# 保存最优出库方案
optimal_Q_A = res.X[:T]  # 深圳水库A最优出库序列
optimal_Q_B = res.X[T:]  # 笋岗滞洪区B最优出库序列
#
# df = (pd.DataFrame({
#     '时段': np.arange(T) + 1,
#     '深圳水库出库(m³/s)': optimal_Q_A,
#     '笋岗滞洪区出库(m³/s)': optimal_Q_B
# }))
# df.to_excel('出库流量分配.xlsx', index=False)

# 提取最优解的索引
optimal_index = np.argmin(res.F)

# 从保存的流量过程中提取最优解的流量过程
optimal_flow_process = problem.flow_processes[optimal_index]

# 创建包含所有水闸流量的DataFrame
flow_process_df = (pd.DataFrame({
    '时段': inflow_data['Date Time'],
    '深圳水库出库流量(m³/s)': optimal_Q_A,
    '笋岗滞洪区出库流量(m³/s)': optimal_Q_B,
    '': '',
    '深圳水库水位(m)': optimal_flow_process['深圳水库水位'],
    '深圳水库库容(万m³)': optimal_flow_process['深圳水库库容'],
    '笋岗滞洪区水位(m)': optimal_flow_process['笋岗滞洪区水位'],
    '笋岗滞洪区库容(万m³)': optimal_flow_process['笋岗滞洪区库容'],
    ' ': '',
    '罗湖站流量(m³/s)': optimal_flow_process['罗湖站流量'],
    '罗湖站水位(m)': optimal_flow_process['罗湖站水位'],
    '布吉河口流量(m³/s)': optimal_flow_process['布吉河口流量'],
    '布吉河口水位(m)': optimal_flow_process['布吉河口水位'],

}))

excel_name = '流量水位过程2.xlsx'
flow_process_df.to_excel(excel_name, index=False)

# 输出提示信息
print("\n调度过程已保存至" + excel_name + '。')

end_time = time.perf_counter()  # 记录结束时间
execution_time = end_time - start_time
print(f"程序运行时间：{execution_time:.6f} 秒")

plt.figure(figsize=(7, 5))
plt.scatter(X[:, 0], X[:, 1], s=30, facecolors='none', edgecolors='r')
plt.xlim(xl[0], xu[0])
plt.ylim(xl[1], xu[1])
plt.title("Design Space")
plt.show()

plt.figure(figsize=(7, 5))
plt.scatter(F[:, 0], F[:, 1], s=30, facecolors='none', edgecolors='blue')
plt.title("Objective Space")
plt.show()